import pickle
from collections import Counter
import numpy as np
import matplotlib.pyplot as plt


def save_model(model, file_name):
    """保存模型"""
    with open(file_name, "wb") as f:
        pickle.dump(model, f)


def load_model(file_name):
    """加载模型"""
    with open(file_name, "rb") as f:
        model = pickle.load(f)
    return model


def flatten_lists(lists):
    flatten_list = []
    for l in lists:
        if type(l) == list:
            flatten_list += l
        else:
            flatten_list.append(l)
    return flatten_list


def merge_maps(dict1, dict2):
    """合并两个word2id或合并两个tag2id"""
    for key in dict2.keys():
        if key not in dict1:
            dict1[key] = len(dict1)
    return dict1


def find_key(dic, value):
    keys = list(dic.keys())
    values = list(dic.values())
    idx = values.index(value)
    key = keys[idx]
    return key


def plot_confusion_matrix(cm, tags, title, save=False):
    """绘制混淆矩阵"""
    cm = np.around(cm/ cm.sum(axis=1)[:, np.newaxis],decimals=2)  # 归一化
    plt.figure(figsize=(15,15), dpi=50)
    plt.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)  # 在特定的窗口上显示图像
    plt.title(title, fontsize=10)  # 图像标题
    plt.colorbar()  # 添加自定义颜色条
    cls_num = len(tags)
    num_local = np.array(range(cls_num))
    plt.xticks(num_local, tags, fontsize=10, rotation=90)  # 将标签印在x轴坐标上
    plt.yticks(num_local, tags, fontsize=10, rotation=0)  # 将标签印在y轴坐标上

    thresh = cm.max() / 2.
    iters = np.reshape([[[i, j] for j in range(cls_num)] for i in range(cls_num)], (cm.size, 2))
    for i, j in iters:
        plt.text(j, i, format(cm[i, j]),
                 horizontalalignment="center",
                 color="white" if cm[i, j] > thresh else "black")

    plt.ylabel('True label', fontsize=10)  # 纵轴标签
    plt.xlabel('Predicted label', fontsize=10)  # 横轴标签
    if save:
        plt.savefig("./imgs/confusion_matrix.png")
    plt.show()


class Metrics(object):
    """评价模型，计算精确率，召回率，F1分数"""
    def __init__(self, golden_tags, predict_tags):

        # [[t1, t2], [t3, t4]...] --> [t1, t2, t3, t4...]
        self.golden_tags = flatten_lists(golden_tags)
        self.predict_tags = flatten_lists(predict_tags)

        # 辅助计算的变量
        self.tagset = set(self.golden_tags)
        self.correct_tags_number = self.count_correct_tags()
        self.predict_tags_counter = Counter(self.predict_tags)
        self.golden_tags_counter = Counter(self.golden_tags)

        # 计算精确率
        self.precision_scores = self.cal_precision()

        # 计算召回率
        self.recall_scores = self.cal_recall()

        # 计算F1分数
        self.f1_scores = self.cal_f1()

    def count_correct_tags(self):
        """计算每种标签预测正确的个数(对应精确率、召回率计算公式上的tp)，用于后面精确率以及召回率的计算"""
        correct_dict = {}
        for gold_tag, predict_tag in zip(self.golden_tags, self.predict_tags):
            if gold_tag == predict_tag:
                if gold_tag not in correct_dict:
                    correct_dict[gold_tag] = 1
                else:
                    correct_dict[gold_tag] += 1

        return correct_dict

    def cal_precision(self):

        precision_scores = {}
        for tag in self.tagset:
            precision_scores[tag] = self.correct_tags_number.get(tag, 0) / \
                                    self.predict_tags_counter[tag]

        return precision_scores

    def cal_recall(self):

        recall_scores = {}
        for tag in self.tagset:
            recall_scores[tag] = self.correct_tags_number.get(tag, 0) / \
                                 self.golden_tags_counter[tag]
        return recall_scores

    def cal_f1(self):
        f1_scores = {}
        for tag in self.tagset:
            p, r = self.precision_scores[tag], self.recall_scores[tag]
            f1_scores[tag] = 2 * p * r / (p + r + 1e-10)  # 加上一个特别小的数，防止分母为0
        return f1_scores

    def _cal_weighted_average(self):

        weighted_average = {}
        total = len(self.golden_tags)

        # 计算weighted precisions:
        weighted_average['precision'] = 0.
        weighted_average['recall'] = 0.
        weighted_average['f1_score'] = 0.
        for tag in self.tagset:
            size = self.golden_tags_counter[tag]
            weighted_average['precision'] += self.precision_scores[tag] * size
            weighted_average['recall'] += self.recall_scores[tag] * size
            weighted_average['f1_score'] += self.f1_scores[tag] * size

        for metric in weighted_average.keys():
            weighted_average[metric] /= total

        return weighted_average

    def report_avg(self):
        avg_metrics = self._cal_weighted_average()
        print('{:>9s}  {:>9.4f} {:>9.4f} {:>9.4f} {:>9}'.format(
            'avg/total',
            avg_metrics['precision'],
            avg_metrics['recall'],
            avg_metrics['f1_score'],
            len(self.golden_tags)
        ))



